#!/usr/bin/env bash

set -e

function usage() {
    echo "$( basename $0 ) [-g | -h | -b | -w <wheel_file>]"
    echo "    重新安装onnxruntime"
    echo "    -g: 安装onnxruntime-gpu, 缺省安装cpu版本"
    echo "    -b: 建造wheel文件"
    echo "    -i: 安装wheel文件"
    echo "    -w: 指定要安装的wheel文件"
    echo "    -h: 输出本信息"

    exit 0
}

function build_wheel() {
    setupFile='../setup.py'
    if [ ! -f "${setupFile}" ]
    then
        echo "Error, not found ${setupFile}"
        exit 1
    fi
    
    # 首先build wheel
    if ${gpu}
    then
        python3 ${setupFile} bdist_wheel --wheel_name_suffix=gpu --cuda_version=12.1
    else
        python3 ${setupFile} bdist_wheel 
    fi
}

function reinstall() {
    # uninstall 
    orts=($(pip list|grep onnxruntime|awk '{print $1}'))
    for p in "${orts[@]}"
    do
        echo "Uninstall $p"
        pip uninstall $p
    done
    
    # install
    if [ -z "${wheelFile}" ]
    then
        if ${gpu}
        then
            # wheel="dist/onnxruntime-gpu-1.15.0-cp310-cp310-linux_x86_64.whl"
            wheel=$( find dist -name 'onnxruntime_gpu-*.whl' )
        else
            # wheel="dist/onnxruntime-1.15.0-cp310-cp310-linux_x86_64.whl"
            wheel=$( find dist -name 'onnxruntime-*.whl' )
        fi
    else
        wheel="${wheelFile}"
    fi
    
    if [ -z "${wheel}" ]
    then
        echo "Not found wheel file in dist"
        exit 1
    fi
    
    if [ ! -f ${wheel} ]
    then
        echo "Error: ${wheel} not found"
        exit 1
    fi
    
    pip install ${wheel}
}

to_build=false
to_install=false
gpu=false
wheelFile=""
while getopts "gbhiw:" args
do
    case "${args}" in 
        g)
            gpu=true
            ;;
        b)
            to_build=true
            ;;
        i)
            to_install=true
            ;;
        w)
            wheelFile="${OPTARG}"
            ;;
        *)
            usage;;
    esac
done

echo ${wheelFile}

if ${to_build}
then
    build_wheel
fi

if ${to_install}
then
    reinstall
fi
